{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "591abe36",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "import nltk\n",
    "import pickle\n",
    "from shutil import copyfile\n",
    "from collections import Counter\n",
    "from PIL import Image\n",
    "from vocabulary import Vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f77bdcff",
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_fl = \"./coco_data/annotations/instances_train2017.json\"\n",
    "with open(obj_fl) as json_file:\n",
    "  object_detections = json.load(json_file)\n",
    "\n",
    "CATEGORY_LIST = [4, 5, 22, 43]\n",
    "\n",
    "COUNT_PER_CATEGORY = 1000\n",
    "\n",
    "category_dict = dict()\n",
    "for category_id in CATEGORY_LIST:\n",
    "  category_dict[category_id] = dict()\n",
    "\n",
    "all_images = dict()\n",
    "filtered_images = set()\n",
    "\n",
    "for annotation in object_detections['annotations']:\n",
    "  category_id = annotation['category_id']\n",
    "  image_id = annotation['image_id']\n",
    "  area = annotation['area']\n",
    "  if category_id in CATEGORY_LIST:\n",
    "    if image_id not in category_dict[category_id]:\n",
    "      category_dict[category_id][image_id] = []\n",
    "  if image_id not in all_images:\n",
    "    all_images[image_id] = dict()\n",
    "  if category_id not in all_images[image_id]:\n",
    "    all_images[image_id][category_id] = area\n",
    "  else:\n",
    "    current_area = all_images[image_id][category_id]\n",
    "    if area > current_area:\n",
    "      all_images[image_id][category_id] = area\n",
    "\n",
    "if COUNT_PER_CATEGORY == -1:\n",
    "  for category_id in category_dict:\n",
    "    print(\"Processing category {}\".format(category_id))\n",
    "    filtered_images.update(category_dict[category_id].keys())\n",
    "    print(\"  Filtered total {} images of category {}\".format(len(category_dict[category_id].keys()), category_id))\n",
    "else:\n",
    "  for image_id in all_images:\n",
    "    areas = list(all_images[image_id].values())\n",
    "    categories = list(all_images[image_id].keys())\n",
    "    sorted_areas = sorted(areas, reverse=True)\n",
    "    sorted_categories = []\n",
    "    for area in sorted_areas:\n",
    "      sorted_categories.append(categories[areas.index(area)])\n",
    "    all_images[image_id] = sorted_categories\n",
    "\n",
    "  for category_id in category_dict:\n",
    "    print(\"Processing category {}\".format(category_id))\n",
    "    for image_id in category_dict[category_id]:\n",
    "      category_dict[category_id][image_id] = all_images[image_id]\n",
    "    prominance_index = 0\n",
    "    prominent_image_ids = []\n",
    "    while len(category_dict[category_id]) > 0 and len(prominent_image_ids) < COUNT_PER_CATEGORY:\n",
    "      remaining_count = COUNT_PER_CATEGORY - len(prominent_image_ids)\n",
    "      image_ids = []\n",
    "      for image_id in category_dict[category_id]:\n",
    "        if category_dict[category_id][image_id].index(category_id) == prominance_index:\n",
    "          image_ids.append(image_id)\n",
    "      for image_id in image_ids:\n",
    "        del category_dict[category_id][image_id]\n",
    "      if len(image_ids) <= remaining_count:\n",
    "        prominent_image_ids = prominent_image_ids + image_ids\n",
    "        if prominance_index > 4:\n",
    "          print(image_ids)\n",
    "        print(\"  Added all {} images at prominance_index {}\".format(len(image_ids), prominance_index))\n",
    "      else:\n",
    "        random.shuffle(image_ids)\n",
    "        prominent_image_ids = prominent_image_ids + image_ids[0:remaining_count]\n",
    "        print(\"  Added {} images at prominance_index {} out of {} images\".format(remaining_count, prominance_index, len(image_ids)))\n",
    "      prominance_index = prominance_index + 1\n",
    "    filtered_images.update(prominent_image_ids)\n",
    "    print(\"  Completed filtering of total {} images of category {}\".format(len(prominent_image_ids), category_id))\n",
    "\n",
    "print(\"Processed all categories. Number of filtered images is {}\".format(len(filtered_images)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bccf4a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "caps_fl = \"./coco_data/annotations/captions_train2017.json\"\n",
    "with open(caps_fl) as json_file:\n",
    "  captions = json.load(json_file)\n",
    "\n",
    "filtered_annotations = []\n",
    "for annotation in captions['annotations']:\n",
    "  if annotation['image_id'] in filtered_images:\n",
    "    filtered_annotations.append(annotation)\n",
    "captions['annotations'] = filtered_annotations\n",
    "print(\"Number of filtered annotations is {}\".format(len(captions['annotations'])))\n",
    "\n",
    "images = []\n",
    "filtered_image_file_names = set()\n",
    "for image in captions['images']:\n",
    "  if image['id'] in filtered_images:\n",
    "    images.append(image)\n",
    "    filtered_image_file_names.add(image['file_name'])\n",
    "captions['images'] = images\n",
    "print(\"Expected number of filtered images is {}, actual number is {}\".format(len(filtered_images), len(captions['images'])))\n",
    "\n",
    "with open(\"./coco_data/captions.json\", 'w+') as output_file:\n",
    "  json.dump(captions, output_file)\n",
    "!rm -rf ./coco_data/annotations\n",
    "\n",
    "!mkdir coco_data/images\n",
    "for file_name in filtered_image_file_names:\n",
    "  copyfile(\"./coco_data/train2017/{}\".format(file_name),\n",
    "           \"./coco_data/images/{}\".format(file_name))\n",
    "!rm -rf ./coco_data/train2017"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b3a6fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "nltk.download('punkt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2a1e655",
   "metadata": {},
   "outputs": [],
   "source": [
    "def resize_images(input_path, output_path, new_size):\n",
    "    if not os.path.exists(output_path):\n",
    "        os.makedirs(output_path)\n",
    "    image_files = os.listdir(input_path)\n",
    "    num_images = len(image_files)\n",
    "    for i, img in enumerate(image_files):\n",
    "        img_full_path = os.path.join(input_path, img)\n",
    "        with open(img_full_path, 'r+b') as f:\n",
    "            with Image.open(f) as image:\n",
    "                image = image.resize(new_size, Image.ANTIALIAS)\n",
    "                img_sv_full_path = os.path.join(output_path, img)\n",
    "                image.save(img_sv_full_path, image.format)\n",
    "        if (i+1) % 100 == 0 or (i+1) == num_images:\n",
    "            print(\"Resized {} out of {} total images.\".format(i+1, num_images))\n",
    "\n",
    "input_path = './coco_data/images/'\n",
    "output_path = './coco_data/resized_images/'\n",
    "new_size = [256, 256]\n",
    "resize_images(input_path, output_path, new_size)\n",
    "\n",
    "!rm -rf ./coco_data/images\n",
    "!mv ./coco_data/resized_images ./coco_data/images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa13b643",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_vocabulary(json_path, threshold):\n",
    "  with open(json_path) as json_file:\n",
    "    captions = json.load(json_file)\n",
    "  counter = Counter()\n",
    "  i = 0\n",
    "  for annotation in captions['annotations']:\n",
    "    i = i + 1\n",
    "    caption = annotation['caption']\n",
    "    tokens = nltk.tokenize.word_tokenize(caption.lower())\n",
    "    counter.update(tokens)\n",
    "    if i % 1000 == 0 or i == len(captions['annotations']):\n",
    "      print(\"Tokenized {} out of total {} captions.\".format(i, len(captions['annotations'])))\n",
    "\n",
    "  tokens = [tkn for tkn, i in counter.items() if i >= threshold]\n",
    "\n",
    "  vocabulary = Vocabulary()\n",
    "  vocabulary.add_token('<pad>')\n",
    "  vocabulary.add_token('<start>')\n",
    "  vocabulary.add_token('<end>')\n",
    "  vocabulary.add_token('<unk>')\n",
    "\n",
    "  for i, token in enumerate(tokens):\n",
    "    vocabulary.add_token(token)\n",
    "  return vocabulary\n",
    "\n",
    "vocabulary = build_vocabulary(json_path='coco_data/captions.json', threshold=4)\n",
    "vocabulary_path = './coco_data/vocabulary.pkl'\n",
    "with open(vocabulary_path, 'wb') as f:\n",
    "    pickle.dump(vocabulary, f)\n",
    "print(\"Total vocabulary size: {}\".format(len(vocabulary)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
